package com.alison.util;

import com.zaxxer.hikari.HikariDataSource;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;

import javax.sql.DataSource;
import java.lang.reflect.Field;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.Enumeration;
import java.util.HashSet;
import java.util.Map;
import java.util.Properties;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @Author alison
 * @Date 2024/3/28 11:09
 * @Version 1.0
 * @Description <br/>
 * 使用
 * <code>
 * DynamicDataSourceUtil.initDataSource(properties)
 * DynamicDataSourceUtil.getDataSource()
 * DynamicDataSourceUtil.getDataSource(key)
 * DynamicDataSourceUtil.registerDataSource(String...)
 * DynamicDataSourceUtil.deleteDataSource(key)
 * </code>
 */
public class DynamicDataSourceUtil {

    private final static String DRIVER_CLASS_NAME = "spring.datasource.default.driverClassName";
    private final static String URL = "spring.datasource.default.url";
    private final static String USERNAME = "spring.datasource.default.username";
    private final static String PASSWORD = "spring.datasource.default.password";
    private static DynamicDataSource dynamicDataSource = new DynamicDataSource();


    public static void main(String[] args) throws Exception {
        Map<String, String> env = System.getenv();
        Properties properties = new Properties();
        env.forEach((k, v) -> {
            properties.put(k, v);
        });
        initDataSource(properties);
        HikariDataSource dataSource = (HikariDataSource) getDataSource();
        System.out.println(dataSource.getJdbcUrl());
        Assert.notNull(dataSource);

        String driverClassName = "com.mysql.cj.jdbc.Driver";
        String jdbcUrl = "jdbc:mysql://localhost:3306/jpa";
        String username = "root";
        String password = "root";

        registerDataSource("slave", driverClassName, jdbcUrl, username, password);
        DataSource dataSource1 = getDataSource("slave");
        Assert.state(jdbcUrl.equals(dataSource1.getConnection().getMetaData().getURL()));
    }

    /*
     加载Properties文件，
     spring.datasource.default.driverClassName=com.mysql.cj.jdbc.Driver
     spring.datasource.default.url=jdbc:mysql://localhost:3306/test1
     spring.datasource.default.username=root
     spring.datasource.default.password=root
     spring.datasource.master.driverClassName=com.mysql.cj.jdbc.Driver
     spring.datasource.master.url=jdbc:mysql://localhost:3306/test1
     spring.datasource.master.username=root
     spring.datasource.master.password=root
      */
    public static void initDataSource(Properties properties) {
        validateConfig(properties);
        DataSource defaultDataSource = genDefaultDataSource(properties);
//        DataSource defaultDataSource = dataSourceFactory.createDefaultDataSource();
        Map targetDatasources = genTargetDatasources(properties);
        dynamicDataSource.setDefaultTargetDataSource(defaultDataSource);
        dynamicDataSource.setTargetDataSources(targetDatasources);
        dynamicDataSource.afterPropertiesSet();
    }

    public static DataSource getDataSource() {
        return getDataSource(null);
    }

    public static DataSource getDataSource(String key) {
        DataSourceKeyHolder.set(key);
        return dynamicDataSource.determineTargetDataSource();
    }

    public static Connection getConnection(String key) throws SQLException {
        DataSourceKeyHolder.set(key);
        return dynamicDataSource.getConnection();
    }

    public static void registerDataSource(String key, String driverClassName, String url, String username, String password) {
        DataSource dataSource = genDataSource(driverClassName, url, username, password);
        // 可以覆盖, getResolvedDataSources 获取到的 immutableMap
//        dynamicDataSource.getResolvedDataSources().put(key, dataSource);
        try {
            Field resolvedDataSourcesFiled = AbstractRoutingDataSource.class.getDeclaredField("resolvedDataSources");
            resolvedDataSourcesFiled.setAccessible(true);
            Map<String, DataSource> resolvedDataSources = (Map<String, DataSource>) resolvedDataSourcesFiled.get(dynamicDataSource);
            resolvedDataSources.put(key, dataSource);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    public static void deleteDataSource(String key) {
        // 可以覆盖, getResolvedDataSources 获取到的 immutableMap
//        dynamicDataSource.getResolvedDataSources().put(key, dataSource);
        try {
            Field resolvedDataSourcesFiled = AbstractRoutingDataSource.class.getDeclaredField("resolvedDataSources");
            resolvedDataSourcesFiled.setAccessible(true);
            Map<String, DataSource> resolvedDataSources = (Map<String, DataSource>) resolvedDataSourcesFiled.get(dynamicDataSource);
            resolvedDataSources.remove(key);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }


    private static Map<String, DataSource> genTargetDatasources(Properties properties) {
        Enumeration<Object> keys = properties.keys();
        Set<String> datasourceKeys = new HashSet<>();
        Map<String, DataSource> targetDataSources = new ConcurrentHashMap<>(10);
        while (keys.hasMoreElements()) {
            String key = (String) keys.nextElement();
            if (key.contains("default") || !key.contains("spring.datasource")) {
                continue;
            }
            int fromIndex = key.indexOf(".", key.indexOf("datasource") + 10);
            String simpleKey = key.substring(fromIndex + 1, key.indexOf(".", fromIndex + 1));
            datasourceKeys.add(simpleKey);
        }
        datasourceKeys.forEach(key -> {
            HikariDataSource targetDataSource = new HikariDataSource();
            properties.keySet().stream().filter(x -> String.valueOf(x).contains(key)).forEach(k -> {
                String driverClassName = properties.getProperty(String.format("spring.datasource.%s.driverClassName", key));
                String url = properties.getProperty(String.format("spring.datasource.%s.url", key));
                String username = properties.getProperty(String.format("spring.datasource.%s.username", key));
                String password = properties.getProperty(String.format("spring.datasource.%s.password", key));
                targetDataSource.setDriverClassName(driverClassName);
                targetDataSource.setJdbcUrl(url);
                targetDataSource.setUsername(username);
                targetDataSource.setPassword(password);
            });
            targetDataSources.put(key, targetDataSource);
        });
        return targetDataSources;
    }

    public static DataSource genDataSource(String key, Properties properties) {
        return null;
    }


    public static DataSource genDefaultDataSource(Properties properties) {
        String driverClassName = properties.getProperty(DRIVER_CLASS_NAME);
        String url = properties.getProperty(URL);
        String username = properties.getProperty(USERNAME);
        String password = properties.getProperty(PASSWORD);
        return genDataSource(driverClassName, url, username, password);
    }

    public static DataSource genDataSource(String driverClassName, String url, String username, String password) {
        HikariDataSource hikariDataSource = new HikariDataSource();
        hikariDataSource.setDriverClassName(driverClassName);
        hikariDataSource.setJdbcUrl(url);
        hikariDataSource.setUsername(username);
        hikariDataSource.setPassword(password);
        return hikariDataSource;
    }


    private static void validateConfig(Properties properties) {
        String driverClassName = properties.getProperty(DRIVER_CLASS_NAME);
        if (driverClassName == null || !StringUtils.hasText(driverClassName)) {
            throw new IllegalArgumentException(DRIVER_CLASS_NAME + " is required.");
        }
        String url = properties.getProperty(URL);
        if (url == null || !StringUtils.hasText(url)) {
            throw new IllegalArgumentException("spring.datasource.default.url is required.");
        }
        String username = properties.getProperty(USERNAME);
        if (url == username || !StringUtils.hasText((String) username)) {
            throw new IllegalArgumentException("spring.datasource.default.username is required.");
        }
        String password = properties.getProperty(PASSWORD);
        if (url.contains("sqlite")) {
            return;
        }
        if (url == password || !StringUtils.hasText(password)) {
            throw new IllegalArgumentException("spring.datasource.default.password is required.");
        }
    }


    static class DynamicDataSource extends AbstractRoutingDataSource {

        public DataSource determineTargetDataSource() {
            return super.determineTargetDataSource();
        }

        @Override
        protected Object determineCurrentLookupKey() {
            return DataSourceKeyHolder.get();
        }
    }

    static class DataSourceKeyHolder {
        private DataSourceKeyHolder() {
        }

        private final static ThreadLocal<String> dataSourceHolder = new ThreadLocal<>();

        static void set(String key) {
            dataSourceHolder.set(key);
        }

        static String get() {
            return dataSourceHolder.get();
        }

        static void clear() {
            dataSourceHolder.remove();
        }
    }
}
